#include <arpa/inet.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <spdlog/spdlog.h>
#include <sys/epoll.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>

#include <csignal>
#include <string>
#include <thread>
#include <mutex>

#include "feat/wave-reader.h"
#include "fstext/fstext-lib.h"
#include "lat/lattice-functions.h"
#include "nnet3/nnet-utils.h"
#include "online2/online-endpoint.h"
#include "online2/online-nnet2-feature-pipeline.h"
#include "online2/online-nnet3-decoding.h"
#include "online2/online-timing.h"
#include "online2/onlinebin-util.h"
#include "util/kaldi-thread.h"

using namespace fst;
using namespace kaldi;

OnlineNnet2FeaturePipelineConfig feature_opts;
nnet3::NnetSimpleLoopedComputationOptions decodable_opts;
LatticeFasterDecoderConfig decoder_opts;
OnlineEndpointConfig endpoint_opts;

BaseFloat chunk_length_secs = 0.18;
BaseFloat output_period = 1;
BaseFloat samp_freq = 16000.0;
int port_num = 5050;
int read_timeout = -1;
bool produce_time = false;

std::string nnet3_rxfilename;
std::string fst_rxfilename;
std::string word_syms_filename;

int sil_chunk_count = 6;

std::mutex g_lock;

// fst::Fst<fst::StdArc> *decode_fst = NULL;
// fst::SymbolTable *word_syms = NULL;

int SetNoBlock(int sock) {
  if (fcntl(sock, F_SETFL, fcntl(sock, F_GETFD, 0) | O_NONBLOCK) == -1) {
    return -1;
  }
  return 0;
}

inline bool SymbolsToNull(const string &s) {
  if (s.find("[FIL]") != std::string::npos ||
      s.find("[SPK]") != std::string::npos ||
      s.find("呃") != std::string::npos ||
      s.find("啊") != std::string::npos) {
    return true;
  }
  return false;
}

namespace kaldi {

class TcpServer {
 public:
  explicit TcpServer();
  ~TcpServer();

  bool Listen(int32 port);
  int32 Accept();

 private:
  struct ::sockaddr_in h_addr_;
  int32 server_desc_;
};

class TcpSocket {
 public:
  explicit TcpSocket(int32 sock);
  ~TcpSocket();

  bool ReadChunk(size_t len, bool &sil);
  Vector<BaseFloat> GetChunk();

  bool Write(const std::string &msg);
  bool WriteLn(const std::string &msg, const std::string &eol = "\n");

  void Disconnect();

 private:
  int32 sock_desc_;
  int16 *samp_buf_;
  size_t buf_len_, has_read_;
};

std::string LatticeToString(const Lattice &lat,
                            const fst::SymbolTable &word_syms) {
  LatticeWeight weight;
  std::vector<int32> alignment;
  std::vector<int32> words;
  GetLinearSymbolSequence(lat, &alignment, &words, &weight);

  std::ostringstream msg("");
  for (size_t i = 0; i < words.size(); i++) {
    if (words[i] == 0) {
      continue;
    }
    std::string s = word_syms.Find(words[i]);
    if (s.empty() || SymbolsToNull(s)) {
      msg << "";
    } else {
      msg << s << " ";
    }
  }
  return msg.str();
}

std::string GetTimeString(int32 t_beg, int32 t_end, BaseFloat time_unit) {
  char buffer[100];
  double t_beg2 = t_beg * time_unit;
  double t_end2 = t_end * time_unit;
  snprintf(buffer, 100, "%.2f %.2f", t_beg2, t_end2);
  return std::string(buffer);
}

int32 GetLatticeTimeSpan(const Lattice &lat) {
  std::vector<int32> times;
  LatticeStateTimes(lat, &times);
  return times.back();
}

std::string LatticeToString(const CompactLattice &clat,
                            const fst::SymbolTable &word_syms) {
  if (clat.NumStates() == 0) {
    SPDLOG_WARN("Empty Lattie.");
    return "";
  }
  CompactLattice best_path_clat;
  CompactLatticeShortestPath(clat, &best_path_clat);

  Lattice best_path_lat;
  ConvertLattice(best_path_clat, &best_path_lat);
  return LatticeToString(best_path_lat, word_syms);
}

void ThreadHandle(int32 client_fd, fst::Fst<fst::StdArc> *decode_fst,
                  fst::SymbolTable *word_syms, int cpu_num,
                  TransitionModel *trans_model, nnet3::AmNnetSimple *am_nnet) {
  cpu_set_t mask;
  CPU_ZERO(&mask);
  cpu_set_t get;
  CPU_SET(cpu_num, &mask);

  if (pthread_setaffinity_np(pthread_self(), sizeof(mask), &mask) < 0)
    SPDLOG_WARN("can not set thread affinity!");

  CPU_ZERO(&get);

  if (pthread_getaffinity_np(pthread_self(), sizeof(get), &get) < 0)
    SPDLOG_WARN("can not get thread affinity!");

  if (CPU_ISSET(cpu_num, &get))
    SPDLOG_INFO("Running on processor {}", sched_getcpu());

  try {
    OnlineNnet2FeaturePipelineInfo feature_info(feature_opts);

    BaseFloat frame_shift = feature_info.FrameShiftInSeconds();
    int32 frame_subsampling = decodable_opts.frame_subsampling_factor;

    SPDLOG_INFO("Loading AM...");
    g_lock.lock();
    nnet3::DecodableNnetSimpleLoopedInfo decodable_info(decodable_opts,
                                                        am_nnet);
    g_lock.unlock();

    TcpSocket socket(client_fd);

    int32 samp_count = 0;  // this is used for output refresh rate
    size_t chunk_len = static_cast<size_t>(chunk_length_secs * samp_freq);
    int32 check_period = static_cast<int32>(samp_freq * output_period);
    int32 check_count = check_period;

    int32 frame_offset = 0;
    int32 sil_count = 0;

    bool eos = false;
    bool sil = false;
    bool pre_sil = false;

    OnlineNnet2FeaturePipeline *feature_pipeline =
        new OnlineNnet2FeaturePipeline(feature_info);
    SingleUtteranceNnet3Decoder *decoder = new SingleUtteranceNnet3Decoder(
        decoder_opts, *trans_model, decodable_info, *decode_fst,
        feature_pipeline);

    while (!eos) {
      if (frame_offset > 20000 || true) {
        frame_offset = 0;
        sil_count = 0;

        pre_sil = false;

        samp_count = 0;
        check_count = check_period;

        delete decoder;
        delete feature_pipeline;

        feature_pipeline = new OnlineNnet2FeaturePipeline(feature_info);
        decoder = new SingleUtteranceNnet3Decoder(decoder_opts, *trans_model,
                                                  decodable_info, *decode_fst,
                                                  feature_pipeline);
      }

      decoder->InitDecoding(frame_offset);
      OnlineSilenceWeighting silence_weighting(
          *trans_model, feature_info.silence_weighting_config,
          decodable_opts.frame_subsampling_factor);
      std::vector<std::pair<int32, BaseFloat>> delta_weights;

      while (true) {
        eos = !socket.ReadChunk(chunk_len, sil);

        if (sil_count == 0) {
          pre_sil = sil;
        }

        sil &&pre_sil ? sil_count++ : sil_count = 0;
        pre_sil = sil;

        if (sil_count >= sil_chunk_count) {
          Vector<BaseFloat> wave_part = socket.GetChunk();
          feature_pipeline->AcceptWaveform(samp_freq, wave_part);
          samp_count += chunk_len;

          feature_pipeline->InputFinished();
          decoder->AdvanceDecoding();
          decoder->FinalizeDecoding();
          frame_offset += decoder->NumFramesDecoded();

          if (decoder->NumFramesDecoded() > 0) {
            CompactLattice lat;
            decoder->GetLattice(true, &lat);
            std::string msg = LatticeToString(lat, *word_syms);

            if (produce_time) {
              int32 t_beg = frame_offset - decoder->NumFramesDecoded();
              int32 t_end = frame_offset;
              msg =
                  GetTimeString(t_beg, t_end, frame_shift * frame_subsampling) +
                  " " + msg;
            }

            if (msg.size() == 4 || msg.size() == 0) {
              // 如果是多字, 或者解码出的内容为空, 不向客户端写入内容
              break;
            }

            SPDLOG_INFO("End of message: {}", msg);
            socket.WriteLn(msg);
            break;
          } else {
            // 如果没有要解码的内容不向客户端写入内容
            break;
          }
        }

        if (eos) {
          feature_pipeline->InputFinished();

          if (silence_weighting.Active() &&
              feature_pipeline->IvectorFeature() != NULL) {
            silence_weighting.ComputeCurrentTraceback(decoder->Decoder());
            silence_weighting.GetDeltaWeights(
                feature_pipeline->NumFramesReady(),
                frame_offset * decodable_opts.frame_subsampling_factor,
                &delta_weights);
            feature_pipeline->UpdateFrameWeights(delta_weights);
          }

          decoder->AdvanceDecoding();
          decoder->FinalizeDecoding();
          frame_offset += decoder->NumFramesDecoded();
          if (decoder->NumFramesDecoded() > 0) {
            CompactLattice lat;
            decoder->GetLattice(true, &lat);
            std::string msg = LatticeToString(lat, *word_syms);

            // get time-span from previous endpoint to end of audio,
            if (produce_time) {
              int32 t_beg = frame_offset - decoder->NumFramesDecoded();
              int32 t_end = frame_offset;
              msg =
                  GetTimeString(t_beg, t_end, frame_shift * frame_subsampling) +
                  " " + msg;
            }

            KALDI_VLOG(1) << "EndOfAudio, sending message: " << msg;
            socket.WriteLn(msg);
          } else
            socket.Write("\n");
          break;
        }

        Vector<BaseFloat> wave_part = socket.GetChunk();
        feature_pipeline->AcceptWaveform(samp_freq, wave_part);
        samp_count += chunk_len;

        if (silence_weighting.Active() &&
            feature_pipeline->IvectorFeature() != NULL) {
          silence_weighting.ComputeCurrentTraceback(decoder->Decoder());
          silence_weighting.GetDeltaWeights(
              feature_pipeline->NumFramesReady(),
              frame_offset * decodable_opts.frame_subsampling_factor,
              &delta_weights);
          feature_pipeline->UpdateFrameWeights(delta_weights);
        }

        decoder->AdvanceDecoding();

        if (samp_count > check_count) {
          if (decoder->NumFramesDecoded() > 0) {
            Lattice lat;
            decoder->GetBestPath(false, &lat);
            TopSort(&lat);  // for LatticeStateTimes(),
            std::string msg = LatticeToString(lat, *word_syms);

            // get time-span after previous endpoint,
            if (produce_time) {
              int32 t_beg = frame_offset;
              int32 t_end = frame_offset + GetLatticeTimeSpan(lat);
              msg =
                  GetTimeString(t_beg, t_end, frame_shift * frame_subsampling) +
                  " " + msg;
            }

            KALDI_VLOG(1) << "Temporary transcript: " << msg;
            socket.WriteLn(msg, "\r");
          }
          check_count += check_period;
        }

        if (decoder->EndpointDetected(endpoint_opts)) {
          decoder->FinalizeDecoding();
          frame_offset += decoder->NumFramesDecoded();
          CompactLattice lat;
          decoder->GetLattice(true, &lat);
          std::string msg = LatticeToString(lat, *word_syms);

          // get time-span between endpoints,
          if (produce_time) {
            int32 t_beg = frame_offset - decoder->NumFramesDecoded();
            int32 t_end = frame_offset;
            msg = GetTimeString(t_beg, t_end, frame_shift * frame_subsampling) +
                  " " + msg;
          }

          if (msg.size() == 0) {
            break;
          }

          KALDI_VLOG(1) << "Endpoint, sending message: " << msg;
          socket.WriteLn(msg);
          break;  // while (true)
        }
      }
    }
  } catch (const std::exception &e) {
    std::cerr << e.what();
  }
}

}  // namespace kaldi

int main(int argc, char **argv) {
  try {
    using namespace kaldi;
    using namespace fst;

    typedef kaldi::int32 int32;
    typedef kaldi::int64 int64;

    spdlog::set_pattern("[%Y-%m-%d %H:%M:%S.%e] [%^%l%$] [%@] %v");

    const char *usage =
        "Reads in audio from a network socket and performs online\n"
        "decoding with neural nets (nnet3 setup), with iVector-based\n"
        "speaker adaptation and endpointing.\n"
        "Note: some configuration values and inputs are set via config\n"
        "files whose filenames are passed as options\n"
        "\n"
        "Usage: online2-tcp-nnet3-decode-faster [options] <nnet3-in> "
        "<fst-in> <word-symbol-table>\n";

    ParseOptions po(usage);

    po.Register(
        "samp-freq", &samp_freq,
        "Sampling frequency of the input signal (coded as 16-bit slinear).");
    po.Register("chunk-length", &chunk_length_secs,
                "Length of chunk size in seconds, that we process.");
    po.Register("output-period", &output_period,
                "How often in seconds, do we check for changes in output.");
    po.Register("num-threads-startup", &g_num_threads,
                "Number of threads used when initializing iVector extractor.");
    po.Register("read-timeout", &read_timeout,
                "Number of seconds of timout for TCP audio data to appear on "
                "the stream. Use -1 for blocking.");
    po.Register("port-num", &port_num,
                "Port number the server will listen on.");
    po.Register("produce-time", &produce_time,
                "Prepend begin/end imes between endpoints (e.g. '5.46 6.81 "
                "<text_output>', in seconds)");
    po.Register("sil-chunk-count", &sil_chunk_count,
                "Numbers of silence chunk, that we process");

    read_timeout = 1000 * read_timeout;

    feature_opts.Register(&po);
    decodable_opts.Register(&po);
    decoder_opts.Register(&po);
    endpoint_opts.Register(&po);

    po.Read(argc, argv);

    if (po.NumArgs() != 3) {
      po.PrintUsage();
      return 1;
    }

    nnet3_rxfilename = po.GetArg(1);
    fst_rxfilename = po.GetArg(2);
    word_syms_filename = po.GetArg(3);

    signal(SIGPIPE, SIG_IGN);

    TcpServer server;
    server.Listen(port_num);

    SPDLOG_INFO("Loading FST...");

    fst::Fst<fst::StdArc> *decode_fst = ReadFstKaldiGeneric(fst_rxfilename);

    fst::SymbolTable *word_syms = NULL;
    if (!word_syms_filename.empty())
      if (!(word_syms = fst::SymbolTable::ReadText(word_syms_filename)))
        KALDI_ERR << "Could not read symbol table from file "
                  << word_syms_filename;


    SPDLOG_INFO("Loading Model..");

    TransitionModel *trans_model = new TransitionModel();
    nnet3::AmNnetSimple *am_nnet = new nnet3::AmNnetSimple();
    {
      bool binary;
      Input ki(nnet3_rxfilename, &binary);
      trans_model->Read(ki.Stream(), binary);
      am_nnet->Read(ki.Stream(), binary);
      SetBatchnormTestMode(true, &(am_nnet->GetNnet()));
      SetDropoutTestMode(true, &(am_nnet->GetNnet()));
      nnet3::CollapseModel(nnet3::CollapseModelConfig(), &(am_nnet->GetNnet()));
    }

    SPDLOG_INFO("Loading Model Finished!");

    int cpu_num = 0;

    while (true) {
      int32 client_desc = server.Accept();
      if (client_desc < 0) {
        SPDLOG_CRITICAL("Accept client failed! Breaking...");
        close(client_desc);
        continue;
      }

      std::thread t = std::thread(ThreadHandle, client_desc, decode_fst,
                                  word_syms, cpu_num, trans_model, am_nnet);
      t.detach();

      if (cpu_num < sysconf(_SC_NPROCESSORS_ONLN) - 1) {
        cpu_num++;
      } else {
        cpu_num = 0;
      }
    }
  } catch (const std::exception &e) {
    std::cerr << e.what();
    return -1;
  }
}

namespace kaldi {
TcpServer::TcpServer() { server_desc_ = -1; }

TcpServer::~TcpServer() {
  if (server_desc_ != -1) {
    close(server_desc_);
  }
}

bool TcpServer::Listen(int32 port) {
  h_addr_.sin_addr.s_addr = INADDR_ANY;
  h_addr_.sin_port = htons(port);
  h_addr_.sin_family = AF_INET;

  server_desc_ = socket(AF_INET, SOCK_STREAM, 0);

  if (server_desc_ == -1) {
    SPDLOG_ERROR("Cannot create TCP socket!");
    return false;
  }

  int32 flag = 1;
  int32 len = sizeof(int32);

  if (setsockopt(server_desc_, SOL_SOCKET, SO_REUSEADDR, &flag, len) == -1) {
    SPDLOG_ERROR("Cannot set socket options!");
    return false;
  }

  if (bind(server_desc_, (struct sockaddr *)&h_addr_, sizeof(h_addr_)) == -1) {
    SPDLOG_ERROR("Cannot bind to port: {} (is it taken?)", port);
    return false;
  }

  if (listen(server_desc_, 1) == -1) {
    SPDLOG_ERROR("Cannot listen on port!");
    return false;
  }

  SPDLOG_INFO("TcpServer: Listening on port: {}", port);
  return true;
}

int32 TcpServer::Accept() {
  SPDLOG_INFO("Waiting for client ...");

  socklen_t len;

  len = sizeof(struct sockaddr);
  struct sockaddr_in client_sockaddr;
  int client_desc =
      accept(server_desc_, (struct sockaddr *)&client_sockaddr, &len);

  struct sockaddr_storage addr;
  char ipstr[20];

  len = sizeof(addr);
  getpeername(client_desc, (struct sockaddr *)&addr, &len);

  struct sockaddr_in *s = (struct sockaddr_in *)&addr;
  inet_ntop(AF_INET, &s->sin_addr, ipstr, sizeof(ipstr));

  SPDLOG_INFO("TcpServer: Accepted connection from: {}", ipstr);
  return client_desc;
}

TcpSocket::TcpSocket(int32 sock) {
  sock_desc_ = sock;
  samp_buf_ = NULL;
  buf_len_ = 0;
}

TcpSocket::~TcpSocket() {
  if (sock_desc_ != -1) {
    close(sock_desc_);
  }
}

bool TcpSocket::ReadChunk(size_t len, bool &sil) {
  if (buf_len_ != len) {
    buf_len_ = len;
    delete[] samp_buf_;
    samp_buf_ = new int16[len];
  }

  ssize_t ret;
  char *samp_buf_p = reinterpret_cast<char *>(samp_buf_);
  size_t to_read = len * sizeof(int16);
  has_read_ = 0;

  int count_zero = 0;
  while (to_read > 0) {
    ret =
        read(sock_desc_, static_cast<void *>(samp_buf_p + has_read_), to_read);
    if (ret <= 0) {
      SPDLOG_WARN("Stream over...");
      break;
    }
    to_read -= ret;
    has_read_ += ret;
  }
  has_read_ /= sizeof(int16);
  for (int i = 0; i < has_read_; i++) {
    if (samp_buf_[i] == 0) {
      count_zero++;
    }
  }

  sil = count_zero >= has_read_ * 0.6 ? true : false;

  return has_read_ > 0;
}

Vector<BaseFloat> TcpSocket::GetChunk() {
  Vector<BaseFloat> buf;

  buf.Resize(static_cast<MatrixIndexT>(has_read_));

  for (int i = 0; i < has_read_; i++) {
    buf(i) = static_cast<BaseFloat>(samp_buf_[i]);
  }

  return buf;
}

bool TcpSocket::Write(const std::string &msg) {
  const char *p = msg.c_str();
  size_t to_write = msg.size();
  size_t wrote = 0;
  while (to_write > 0) {
    ssize_t ret =
        write(sock_desc_, static_cast<const void *>(p + wrote), to_write);
    if (ret <= 0) return false;

    to_write -= ret;
    wrote += ret;
  }
  int flag = 1;
  setsockopt(sock_desc_, IPPROTO_TCP, TCP_NODELAY, (char *)&flag, sizeof(int));
  return true;
}

bool TcpSocket::WriteLn(const std::string &msg, const std::string &eol) {
  if (Write(msg))
    return Write(eol);
  else
    return false;
}
}  // namespace kaldi
